use crate::crypto::{hmac_sha3, KEY_SIZE};
use crate::types::*;
use crate::util::*;

use num_traits::identities::{One, Zero};

use aes::cipher::generic_array::GenericArray;
use aes::cipher::KeySizeUser;
use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use aes::Aes256;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha20Rng;
type Aes256CbcEnc = cbc::Encryptor<Aes256>;
type Aes256CbcDec = cbc::Decryptor<Aes256>;

// This submodule contains the code for (de)packetizer and crypto,
// including trial decode

/**
 * Packet format
 */

#[derive(Clone, Debug)]
pub struct Header {
    pub msg_size: u16,
    pub stream_id: Vec<u8>,
    pub sequence_no: usize,
    pub enc_hdr_as_vec: Vec<u8>,
}

#[derive(Clone, Debug)]
pub struct Packet {
    pub header: Header,
    pub payload: Vec<u8>,
}

/***********************
 * MESSAGE <-> PACKETS *
 ***********************/

pub struct Packetizer<'a> {
    pub rng: rand_chacha::ChaCha20Rng,
    pub aes_key: GenericArray<u8, <aes::Aes256 as KeySizeUser>::KeySize>,
    pub messages_iter: &'a mut dyn Iterator<Item = Vec<u8>>,
    pub stream_id: &'a [u8; 8],
    /// only 6 bytes are used
    pub sequence_no: usize,
}

impl<'a> Packetizer<'a> {
    pub fn with_live_random(
        key: &[u8],
        message_chunks_iter: &'a mut dyn Iterator<Item = Vec<u8>>,
        stream_id: &'a [u8; 8],
    ) -> Self {
        Packetizer::new(
            key,
            message_chunks_iter,
            stream_id,
            ChaCha20Rng::from_rng(rand::thread_rng()).unwrap(),
        )
    }

    pub fn new(
        key: &[u8],
        message_chunks_iter: &'a mut dyn Iterator<Item = Vec<u8>>,
        stream_id: &'a [u8; 8],
        rng: rand_chacha::ChaCha20Rng,
    ) -> Self {
        Packetizer {
            rng,
            aes_key: *GenericArray::from_slice(key),
            messages_iter: message_chunks_iter,
            stream_id,
            sequence_no: 0,
        }
    }
}

/// Message format:
/// - header:
///   - random iv1 (16 bytes)
///   - AES256Cbc(key, iv1, nonce) (16 bytes)
///   - header proper, encrypted by xor'ing with nonce
///     - stream_id       (8 bytes)
///     - ciphertext size (u16 = 2 bytes)
///     - sequence no.    (6 bytes)
///   - HMAC (encrypted header) truncated to 16 bytes
///   - HMAC (iv1 || encrypted nonce || encrypted header proper || iv2 || ciphertext), truncated to 16 bytes
///     i.e. including all data from payload
/// - payload:
///   - random iv2 (16 bytes)
///   - ciphertext, consisting of message encrypted as
///     AES256Cbc(key, iv2, message)
impl<'a> Iterator for Packetizer<'a> {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        match self.messages_iter.next() {
            Some(message) => {
                let key = self.aes_key;

                let mut iv1 = vec![0u8; 16];
                self.rng.fill(&mut iv1[..]);
                let iv1 = *GenericArray::from_slice(&iv1);

                let mut nonce_vec = vec![0u8; 16];
                self.rng.fill(&mut nonce_vec[..]);

                let cipher_n = Aes256CbcEnc::new(&key, &iv1);

                let encrypted_nonce = cipher_n.encrypt_padded_vec_mut::<Pkcs7>(&nonce_vec[..]);

                let mut iv2 = vec![0u8; 16];
                self.rng.fill(&mut iv2[..]);
                let iv2 = *GenericArray::from_slice(&iv2);

                let cipher = Aes256CbcEnc::new(&key, &iv2);
                let ciphertext = cipher.encrypt_padded_vec_mut::<Pkcs7>(&message[..]);

                let mut hdr = vec![];
                // 2 bytes are plenty, considering the inefficiency of the encoding.
                let payload_size: usize = ciphertext.len() + iv2.len();

                let payload_size: u16 = u16::try_from(payload_size).ok()?;
                // TODO refuse with error message if this overflows u16, of course:
                // ok_or(anyhow::Error::msg(format!("{} too large for u16", size)))?;
                let mut payload_size_be = u16::to_be_bytes(payload_size).to_vec();

                // 8 bytes for stream id
                hdr.extend_from_slice(self.stream_id);
                // 2 bytes for size of message
                hdr.append(&mut payload_size_be);
                // 6 bytes for sequence number
                hdr.extend_from_slice(&usize::to_be_bytes(self.sequence_no)[2..8]);

                self.sequence_no += 1;

                let mut xor_hdr = vec![];
                for i in 0..hdr.len() {
                    xor_hdr.push(hdr[i] ^ nonce_vec[i]);
                }

                let mut hmac_hdr: Vec<u8> = hmac_sha3(&key[..], &xor_hdr[..]);
                let mut enc_hdr_msg = vec![];
                enc_hdr_msg.extend_from_slice(&iv1[..]);
                enc_hdr_msg.extend_from_slice(&encrypted_nonce[..]);
                enc_hdr_msg.extend_from_slice(&xor_hdr[..]);
                enc_hdr_msg.extend_from_slice(&iv2[..]);
                enc_hdr_msg.extend_from_slice(&ciphertext[..]);
                let mut hmac_hdr_msg: Vec<u8> = hmac_sha3(&key[..], &enc_hdr_msg[..]);

                hmac_hdr.truncate(16);
                hmac_hdr_msg.truncate(16);

                let mut enc_msg = vec![];
                enc_msg.append(&mut iv1.to_vec());
                enc_msg.append(&mut encrypted_nonce.to_vec()); // length: two blocks
                enc_msg.append(&mut xor_hdr);
                enc_msg.append(&mut hmac_hdr);
                enc_msg.append(&mut hmac_hdr_msg);
                enc_msg.append(&mut iv2.to_vec());
                enc_msg.append(&mut ciphertext.to_vec());

                Some(enc_msg)
            }
            None => None,
        }
    }
}

/****************
 * Depacketizer *
 ****************/

#[derive(Clone, Debug)]
pub enum DepacketizeError {
    HeaderHmacMismatch,
    HeaderDecryptError,
    HeaderAndMessageHmacMismatch,
    DecryptError,
    TooShortForHeader,
    TooShortForHeaderAndMessageHmac,
    TooShortForMessage,
}

// TODO no magic numbers, make explicit how this value is computed
pub static HEADER_SIZE: u16 = 80;

/// See if it is already possible to glean a packet header from the given input,
/// which may just be a prefix of the full input.
/// Intended to be used to detect the start of a packet, to "synchronize" decoder
/// with message stream.
pub fn try_read_packet_header(key: &[u8], input: &[u8]) -> Result<Header, DepacketizeError> {
    // TODO header size and layout should be clearly declared, not just hardcoded
    // in this function and its counterpart in the packetizer.
    if input.len() < HEADER_SIZE as usize {
        return Err(DepacketizeError::TooShortForHeader);
    }

    let (iv1, input) = input.split_at(16);
    let (encrypted_nonce, input) = input.split_at(32); // two blocks of Cbc
    let (xor_hdr, input) = input.split_at(16);
    let (hmac_hdr, _) = input.split_at(16);

    let my_hmac_hdr: Vec<u8> = hmac_sha3(key, xor_hdr);
    if hmac_hdr != &my_hmac_hdr[..16] {
        return Err(DepacketizeError::HeaderHmacMismatch);
    }

    let key = *GenericArray::from_slice(&key[0..*KEY_SIZE]);
    let iv1 = *GenericArray::from_slice(iv1);

    let cipher = Aes256CbcDec::new(&key, &iv1);
    let unencrypted_nonce = cipher
        .decrypt_padded_vec_mut::<Pkcs7>(encrypted_nonce)
        .map_err(|_e| DepacketizeError::HeaderDecryptError)?;

    let hdr = xor_blocks(&unencrypted_nonce, xor_hdr);

    let hdrx = &hdr[..];
    let (stream_id, hdrx) = hdrx.split_at(8);
    let (size, seq_no) = hdrx.split_at(2);

    let mut seq_no8: [u8; 8] = Default::default();
    seq_no8[2..].copy_from_slice(seq_no);
    let seq_no = usize::from_be_bytes(seq_no8);

    let mut size2: [u8; 2] = Default::default();
    size2.copy_from_slice(size);
    let size = u16::from_be_bytes(size2);

    let mut enc_hdr_as_vec = vec![];
    enc_hdr_as_vec.extend_from_slice(&iv1[..]);
    enc_hdr_as_vec.extend_from_slice(encrypted_nonce);
    enc_hdr_as_vec.extend_from_slice(xor_hdr);

    Ok(Header {
        msg_size: size,
        stream_id: stream_id.to_vec(),
        sequence_no: seq_no,
        enc_hdr_as_vec: enc_hdr_as_vec.to_vec(),
    })
}

pub fn try_read_packet(key: &[u8], input: &[u8]) -> Result<Packet, DepacketizeError> {
    match try_read_packet_header(key, input) {
        Ok(header) => {
            let size = header.msg_size;

            if input.len() < HEADER_SIZE as usize + 16 {
                return Err(DepacketizeError::TooShortForHeaderAndMessageHmac);
            }

            let (_, input) = input.split_at(HEADER_SIZE as usize);
            let (hmac_hdr_msg, input) = input.split_at(16);

            if input.len() < size.into() {
                return Err(DepacketizeError::TooShortForMessage);
            }

            let (iv, ciphertext) = input.split_at(16);

            let mut enc_hdr_msg = vec![];
            enc_hdr_msg.extend_from_slice(&header.enc_hdr_as_vec[..]);
            enc_hdr_msg.extend_from_slice(iv);
            enc_hdr_msg.extend_from_slice(&ciphertext[..size as usize - iv.len()]);
            let my_hmac_hdr_msg: Vec<u8> = hmac_sha3(key, &enc_hdr_msg[..]);

            if hmac_hdr_msg != &my_hmac_hdr_msg[..16] {
                return Err(DepacketizeError::HeaderAndMessageHmacMismatch); // isn't it?
            }

            let key = *GenericArray::from_slice(key);
            let iv = *GenericArray::from_slice(iv);

            let cipher = Aes256CbcDec::new(&key, &iv);

            let decrypted_ciphertext = cipher
                .decrypt_padded_vec_mut::<Pkcs7>(ciphertext)
                .map_err(|_e| DepacketizeError::DecryptError)?;

            Ok(Packet {
                payload: decrypted_ciphertext.to_vec(),
                header,
            })
        }
        Err(e) => Err(e),
    }
}

trait FnDecode = Fn(Vec<u8>) -> Result<BoundsLP, anyhow::Error>;
trait FnTrialDecodeHeader =
    Fn(&[BoundsLP]) -> Result<Result<Header, DepacketizeError>, anyhow::Error>;
trait FnTrialDecode =
    Fn(u16, &[BoundsLP]) -> Result<Result<Packet, DepacketizeError>, anyhow::Error>;

/// For the general case of externally delimited messages,
/// that does not rely on prefix free property of grammar
/// to separate one received sentence from the next.
pub struct DecodingLoop<'a> {
    fragments_iter: Box<dyn Iterator<Item = BoundsLP> + 'a>,
    collected: Vec<BoundsLP>,
    /// None means "waiting for header".
    awaiting_packet_of_size: Option<u16>,
    fn_trial_decode_header: Box<dyn FnTrialDecodeHeader + 'a>,
    fn_trial_decode: Box<dyn FnTrialDecode + 'a>,
}

impl<'a> DecodingLoop<'a> {
    pub fn new(
        fragments_iter: Box<dyn Iterator<Item = BoundsLP> + 'a>,
        fn_trial_decode_header: Box<dyn FnTrialDecodeHeader + 'a>,
        fn_trial_decode: Box<dyn FnTrialDecode + 'a>,
    ) -> Self {
        DecodingLoop {
            fragments_iter,
            collected: vec![],
            awaiting_packet_of_size: None,
            fn_trial_decode_header,
            fn_trial_decode,
        }
    }
}

impl<'a> Iterator for DecodingLoop<'a> {
    type Item = Packet; // How and when to record errors?

    fn next(&mut self) -> Option<Self::Item> {
        // TODO always record at least the last error upon returning None?
        // TODO (important optimization) avoid duplicate decoding work as much as we can
        //      (cf. internally delimited variant).

        'trial_decoding: loop {
            if self.collected.is_empty() {
                match self.fragments_iter.next() {
                    None => return None,
                    Some(lp) => {
                        self.collected.push(lp);
                        // otherwise drop silently?
                    }
                }
            }

            match self.awaiting_packet_of_size {
                None => {
                    match (self.fn_trial_decode_header)(&mut self.collected) {
                        Ok(Ok(header)) => {
                            self.awaiting_packet_of_size = Some(header.msg_size);
                            // the packet could theoretically already be complete
                            continue 'trial_decoding;
                        }
                        Ok(Err(DepacketizeError::TooShortForHeader)) => {
                            match self.fragments_iter.next() {
                                None => return None,
                                Some(lp) => {
                                    self.collected.push(lp);
                                    // otherwise drop silently?
                                }
                            }
                            continue 'trial_decoding;
                        }
                        Ok(Err(DepacketizeError::HeaderHmacMismatch)) => {
                            // this cannot be the start of a valid packet.
                            self.collected.remove(0);
                            continue 'trial_decoding;
                        }
                        Ok(_) => {
                            unreachable!();
                        }
                        // TODO: error reporting
                        Err(_e) => {
                            // could not decode at all, retry without first part (which is always the offending part)
                            self.collected.remove(0);
                            continue 'trial_decoding;
                        }
                    }
                }
                Some(expected_size) => {
                    match (self.fn_trial_decode)(expected_size, &mut self.collected) {
                        Ok(Ok(packet)) => {
                            // Throw away the stuff at the end, and get ready for the next packet!
                            self.collected = vec![];
                            self.awaiting_packet_of_size = None;
                            // Success!
                            return Some(packet);
                        }
                        Ok(Err(DepacketizeError::TooShortForHeaderAndMessageHmac)) => {
                            // wait for more.
                            match self.fragments_iter.next() {
                                None => return None,
                                Some(lp) => {
                                    self.collected.push(lp);
                                    // otherwise drop silently?
                                }
                            }
                            continue 'trial_decoding;
                        }
                        Ok(Err(DepacketizeError::TooShortForMessage)) => {
                            // wait for more.
                            match self.fragments_iter.next() {
                                None => return None,
                                Some(lp) => {
                                    self.collected.push(lp);
                                    // otherwise drop silently?
                                }
                            }
                            continue 'trial_decoding;
                        }
                        Ok(_x) => {
                            self.collected.remove(0);
                            self.awaiting_packet_of_size = None;
                            continue 'trial_decoding;
                        }
                        Err(_e) => {
                            self.collected.remove(0);
                            self.awaiting_packet_of_size = None;
                            continue 'trial_decoding;
                        }
                    }
                }
            }
        }
    }
}

pub fn combine_fragments(collected_fragments: &[BoundsLP]) -> BoundsLP {
    let mut result = BoundsLP {
        l: Zero::zero(),
        p: One::one(),
    };
    for bounds_lp in collected_fragments.iter() {
        let rl = result.l.clone();
        let rp = result.p.clone();
        result.l = rl.clone() + bounds_lp.l.clone() * rp.clone();
        result.p = bounds_lp.p.clone() * rp.clone();
    }
    result
}

pub fn trial_decode_header_from_fragments(
    key: &[u8],
    collected_fragments: &[BoundsLP],
) -> Result<Result<Header, DepacketizeError>, anyhow::Error> {
    let result = combine_fragments(collected_fragments);
    let input_hyp = binaryfraction_between_0_and_1_to_bytes(&result.l);

    let upper = result.l.clone() + result.p.clone();
    let knownbits = binaryfractions_length_common_prefix(&result.l, &upper);

    let knownbytes = match knownbits {
        BinaryFractionLongestCommonPrefix::Same => todo!(),
        BinaryFractionLongestCommonPrefix::Invalid => todo!(),
        BinaryFractionLongestCommonPrefix::Length { bits } => bits / 8,
    } as usize;

    Ok(try_read_packet_header(
        &key[0..*KEY_SIZE],
        &input_hyp[0..std::cmp::min(input_hyp.len(), knownbytes)],
    ))
}

pub fn trial_decode_from_fragments(
    key: &[u8],
    expected_size: u16,
    collected_fragments: &[BoundsLP],
) -> Result<Result<Packet, DepacketizeError>, anyhow::Error> {
    let result = combine_fragments(collected_fragments);

    let needed = (expected_size + HEADER_SIZE + 16) as usize; // in bytes

    let input_hyp = binaryfraction_between_0_and_1_to_bytes(&result.l);
    let u = result.l.clone() + result.p.clone();

    if !binaryfractions_same_prefix((needed * 8) as u32, &u, &result.l) {
        Ok(Err(DepacketizeError::TooShortForMessage))
    } else {
        Ok(try_read_packet(
            &key[0..*KEY_SIZE],
            &input_hyp[0..(needed) as usize],
        ))
    }
}
